import torch
import torch.utils
import torch.utils.data
import torch.utils.data.dataloader
import torchvision
import numpy as np
from torch import nn
import matplotlib.pyplot as plt
import torchsummary

'''
功能：
    加载数据集函数
参数:
    root:       数据集路径
    batch_size: 批大小
    transforms: 一系列变换
    ratio:      所取数据集的占比
返回:
    总数据条数,训练集与测试集的dataloader对象
'''
def load_data_FMNIST(root, batch_size, transforms_, ratio=1):
    train_data = torchvision.datasets.FashionMNIST(root, train=True, transform=transforms_, download=True)
    test_data = torchvision.datasets.FashionMNIST(root, train=False, transform=transforms_, download=True)
    train_iter = None
    test_iter = None

    if ratio==1:
        train_iter = torch.utils.data.DataLoader(train_data, batch_size, shuffle=True)
        test_iter = torch.utils.data.DataLoader(test_data, batch_size, shuffle=True)
    elif ratio<1 and ratio>0:
        # choice函数：生成指定范围内的下标随机数组
        # 参数1:选择的范围，参数2：选择的个数，参数3：是否可重复
        train_num, test_num = int(len(train_data)*ratio), int(len(test_data)*ratio)
        subset_indices_train = np.random.choice(len(train_data), train_num, False)
        subset_indices_test = np.random.choice(len(test_data), test_num, False)
        train_subset = torch.utils.data.Subset(train_data, subset_indices_train)
        test_subset = torch.utils.data.Subset(test_data, subset_indices_test)
        train_iter = torch.utils.data.DataLoader(train_subset, batch_size, shuffle=True)
        test_iter = torch.utils.data.DataLoader(test_subset, batch_size, shuffle=True)
    else:
        print('Warning: Ratio of dataset must be below 1')

    return (len(train_data)*ratio, len(test_data)*ratio), train_iter, test_iter

'''
功能：
    计算分类问题的每个epoch的准确度
参数：
    net:        模型对象
    data_iter:  DataLoader对象, 测试数据集
返回值：
    准确度
'''
def accuracy(net, data_iter):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    accuracy = 0
    data_num = 0
    for x, y in data_iter:
        x, y = x.to(device), y.to(device)
        out = net(x)
        max_indices = torch.argmax(out, dim=1)
        same = (max_indices==y).float()
        accuracy += same.sum().item()
        data_num += x.shape[0]
    return accuracy/data_num

'''
功能：
    模型初始化
参数：
    module: 模型对象
'''
def module_init(module):
    if type(module)==nn.Linear or type(module)==nn.Conv2d:
        nn.init.xavier_normal_(module.weight)

'''
功能：
    训练函数
参数：
    net:            模型对象
    train_iter:     DataLoader对象, 训练数据集
    test_iter:      DataLoader对象, 测试数据集
    epochs:         epoch数
    optimizer:      优化算法
    loss:           损失函数
    draw:           是否需要绘图
'''
def train(net, train_iter, test_iter, epochs, optimizer, loss, draw):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    net.to(device)
    net.apply(module_init)

    losses = []
    acc = []
    for epoch in range(epochs):
        loss_per_epoch = 0
        batch_num = 0
        for x,y in train_iter:
            x, y = x.to(device), y.to(device)
            l = loss(net(x), y)
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            loss_per_epoch += l.item()
            batch_num += 1
        loss_per_epoch /= batch_num
        losses.append(loss_per_epoch)
        
        acc_per_epoch = accuracy(net, test_iter)
        acc.append(acc_per_epoch)
        print(f'epoch:{epoch+1}, loss:{loss_per_epoch}, accuracy:{acc_per_epoch}')
    
    if draw is True:
        plt.figure()
        plt.plot(losses, color='red', label='loss')
        plt.plot(acc, color='gray', label='accuracy')
        plt.legend()
        plt.show()


if __name__=='__main__':
    num, train, test = load_data_FMNIST('./data', 64, None, 0.1)
    print(num, len(train), len(test))
